#!/usr/bin/env python3
"""
AI Historical Simulation Platform - Startup Script

This script handles system initialization, dependency checking, and provides
an easy entry point for running the AI Historical Simulation Platform.

Features:
- Dependency validation and installation guidance
- System requirements checking
- Graceful error handling and recovery
- Multiple startup options
- Environment validation

Usage:
    python run.py                  # Interactive mode selection
    python run.py --chat          # Direct to chat mode
    python run.py --demo          # Direct to demo mode
    python run.py --check         # Check system requirements
    python run.py --install-deps  # Install missing dependencies
"""

import sys
import os
import subprocess
import platform
import importlib
import json
from pathlib import Path
from typing import List, Dict, Tuple, Optional
import argparse

# Required Python version
MIN_PYTHON_VERSION = (3, 8)
RECOMMENDED_PYTHON_VERSION = (3, 9)

# Required packages with versions
REQUIRED_PACKAGES = {
    'numpy': '>=1.19.0',
    'scipy': '>=1.7.0',
    'scikit-learn': '>=1.0.0',
    'psutil': '>=5.8.0',
    'asyncio': None,  # Built-in
    'pathlib': None,  # Built-in
    'dataclasses': None,  # Built-in in Python 3.7+
    'typing': None,  # Built-in
    'enum': None,  # Built-in
    'logging': None,  # Built-in
    'json': None,  # Built-in
    'time': None,  # Built-in
    'threading': None,  # Built-in
    'hashlib': None,  # Built-in
    'collections': None,  # Built-in
}

# Optional packages for enhanced functionality
OPTIONAL_PACKAGES = {
    'numba': '>=0.56.0',  # For HDC performance acceleration
    'yaml': '>=5.4.0',    # For YAML config files
    'toml': '>=0.10.0',   # For TOML config files
    'aiohttp': '>=3.8.0', # For web interface
    'fastapi': '>=0.70.0', # For API server
    'uvicorn': '>=0.15.0', # For web server
    'openai': '>=0.27.0',  # For OpenAI API integration
    'transformers': '>=4.20.0', # For local LLM models
    'torch': '>=1.12.0',   # For neural network components
}


class Colors:
    """Console color codes for better output formatting."""
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    MAGENTA = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    END = '\033[0m'
    
    @classmethod
    def disable(cls):
        """Disable colors for non-terminal environments."""
        cls.RED = cls.GREEN = cls.YELLOW = cls.BLUE = ''
        cls.MAGENTA = cls.CYAN = cls.WHITE = cls.BOLD = ''
        cls.UNDERLINE = cls.END = ''


# Disable colors if not in terminal
if not sys.stdout.isatty():
    Colors.disable()


def print_header():
    """Print application header."""
    print(f"{Colors.CYAN}{Colors.BOLD}")
    print("=" * 80)
    print("🎭 AI Historical Simulation Platform")
    print("让有趣的灵魂跨越时空在当下聚首")
    print("Let fascinating souls gather across time in the present")
    print("=" * 80)
    print(f"{Colors.END}")


def print_status(message: str, status: str = "info"):
    """Print status message with color coding."""
    if status == "success":
        print(f"{Colors.GREEN}✅ {message}{Colors.END}")
    elif status == "error":
        print(f"{Colors.RED}❌ {message}{Colors.END}")
    elif status == "warning":
        print(f"{Colors.YELLOW}⚠️  {message}{Colors.END}")
    elif status == "info":
        print(f"{Colors.BLUE}ℹ️  {message}{Colors.END}")
    elif status == "working":
        print(f"{Colors.MAGENTA}🔄 {message}{Colors.END}")
    else:
        print(f"{message}")


def check_python_version() -> bool:
    """Check if Python version meets requirements."""
    current_version = sys.version_info[:2]
    
    if current_version < MIN_PYTHON_VERSION:
        print_status(
            f"Python {MIN_PYTHON_VERSION[0]}.{MIN_PYTHON_VERSION[1]}+ required, "
            f"but {current_version[0]}.{current_version[1]} found",
            "error"
        )
        return False
    elif current_version < RECOMMENDED_PYTHON_VERSION:
        print_status(
            f"Python {current_version[0]}.{current_version[1]} detected. "
            f"Python {RECOMMENDED_PYTHON_VERSION[0]}.{RECOMMENDED_PYTHON_VERSION[1]}+ recommended",
            "warning"
        )
    else:
        print_status(f"Python {current_version[0]}.{current_version[1]} ✓", "success")
    
    return True


def check_package_availability(package_name: str, version_req: Optional[str] = None) -> Tuple[bool, Optional[str]]:
    """Check if a package is available and meets version requirements."""
    try:
        module = importlib.import_module(package_name)
        
        if version_req is None:
            return True, None
        
        # Try to get version
        version = None
        for attr in ['__version__', 'version', 'VERSION']:
            if hasattr(module, attr):
                version = getattr(module, attr)
                break
        
        if version is None:
            return True, "unknown"
        
        # Simple version checking (could be improved with packaging library)
        if version_req.startswith('>='):
            required_version = version_req[2:]
            # Simple string comparison (works for most cases)
            return version >= required_version, version
        
        return True, version
        
    except ImportError:
        return False, None


def check_dependencies() -> Dict[str, Dict]:
    """Check all dependencies and return status report."""
    print_status("Checking dependencies...", "working")
    
    results = {
        'required': {},
        'optional': {},
        'missing_required': [],
        'missing_optional': [],
        'all_required_available': True
    }
    
    # Check required packages
    for package, version_req in REQUIRED_PACKAGES.items():
        available, version = check_package_availability(package, version_req)
        results['required'][package] = {
            'available': available,
            'version': version,
            'requirement': version_req
        }
        
        if available:
            version_str = f" (v{version})" if version and version != "unknown" else ""
            print_status(f"{package}{version_str} ✓", "success")
        else:
            print_status(f"{package} - MISSING", "error")
            results['missing_required'].append(package)
            results['all_required_available'] = False
    
    print()
    
    # Check optional packages
    print_status("Checking optional packages...", "working")
    for package, version_req in OPTIONAL_PACKAGES.items():
        available, version = check_package_availability(package, version_req)
        results['optional'][package] = {
            'available': available,
            'version': version,
            'requirement': version_req
        }
        
        if available:
            version_str = f" (v{version})" if version and version != "unknown" else ""
            print_status(f"{package}{version_str} ✓", "success")
        else:
            print_status(f"{package} - not installed", "warning")
            results['missing_optional'].append(package)
    
    return results


def check_system_resources() -> Dict[str, bool]:
    """Check system resource availability."""
    print_status("Checking system resources...", "working")
    
    try:
        import psutil
        
        # Memory check
        memory = psutil.virtual_memory()
        memory_gb = memory.total / (1024**3)
        memory_available_gb = memory.available / (1024**3)
        
        # CPU check
        cpu_count = psutil.cpu_count(logical=True)
        
        # Disk space check
        disk = psutil.disk_usage('/')
        disk_free_gb = disk.free / (1024**3)
        
        print_status(f"Total RAM: {memory_gb:.1f} GB", "info")
        print_status(f"Available RAM: {memory_available_gb:.1f} GB", "info")
        print_status(f"CPU cores: {cpu_count}", "info")
        print_status(f"Free disk space: {disk_free_gb:.1f} GB", "info")
        
        # Check minimum requirements
        requirements_met = {
            'memory': memory_available_gb >= 2.0,  # 2GB minimum
            'cpu': cpu_count >= 2,                # 2 cores minimum
            'disk': disk_free_gb >= 1.0           # 1GB free space
        }
        
        if not requirements_met['memory']:
            print_status("Insufficient memory (2GB+ recommended)", "warning")
        
        if not requirements_met['cpu']:
            print_status("Limited CPU cores (2+ recommended)", "warning")
        
        if not requirements_met['disk']:
            print_status("Low disk space (1GB+ recommended)", "warning")
        
        return requirements_met
        
    except ImportError:
        print_status("Cannot check system resources (psutil not available)", "warning")
        return {'memory': True, 'cpu': True, 'disk': True}


def install_package(package_name: str, version_req: Optional[str] = None) -> bool:
    """Install a package using pip."""
    try:
        package_spec = f"{package_name}{version_req}" if version_req else package_name
        print_status(f"Installing {package_spec}...", "working")
        
        result = subprocess.run(
            [sys.executable, "-m", "pip", "install", package_spec],
            capture_output=True,
            text=True,
            check=True
        )
        
        print_status(f"Successfully installed {package_name}", "success")
        return True
        
    except subprocess.CalledProcessError as e:
        print_status(f"Failed to install {package_name}: {e}", "error")
        return False


def install_missing_dependencies(dependency_results: Dict) -> bool:
    """Install missing required dependencies."""
    missing_required = dependency_results['missing_required']
    
    if not missing_required:
        print_status("All required dependencies are available", "success")
        return True
    
    print_status(f"Installing {len(missing_required)} missing required packages...", "working")
    
    success_count = 0
    for package in missing_required:
        version_req = REQUIRED_PACKAGES[package]
        if install_package(package, version_req):
            success_count += 1
    
    if success_count == len(missing_required):
        print_status("All missing dependencies installed successfully", "success")
        return True
    else:
        print_status(
            f"Failed to install {len(missing_required) - success_count} packages", 
            "error"
        )
        return False


def check_file_structure() -> bool:
    """Check if required files and directories exist."""
    print_status("Checking file structure...", "working")
    
    required_files = [
        'app.py',
        'src/__init__.py',
        'src/hdc/__init__.py',
        'src/hdc/core.py',
        'src/config/__init__.py',
        'src/config/settings.py',
        'src/platform/__init__.py',
        'src/platform/historical_simulation_platform.py',
        'src/personality/__init__.py',
        'src/personality/models.py',
        'src/memory/__init__.py',
        'src/memory/manager.py'
    ]
    
    missing_files = []
    for file_path in required_files:
        if not Path(file_path).exists():
            missing_files.append(file_path)
            print_status(f"Missing: {file_path}", "error")
        else:
            print_status(f"Found: {file_path} ✓", "success")
    
    if missing_files:
        print_status(f"{len(missing_files)} required files are missing", "error")
        return False
    else:
        print_status("All required files found", "success")
        return True


def run_system_check() -> bool:
    """Run comprehensive system check."""
    print_header()
    print_status("Running system compatibility check...", "working")
    print()
    
    # Check Python version
    if not check_python_version():
        return False
    print()
    
    # Check file structure
    if not check_file_structure():
        print_status("Critical files missing. Please ensure complete installation.", "error")
        return False
    print()
    
    # Check dependencies
    dep_results = check_dependencies()
    if not dep_results['all_required_available']:
        print_status("Required dependencies missing", "error")
        return False
    print()
    
    # Check system resources
    resource_check = check_system_resources()
    print()
    
    # Summary
    if dep_results['all_required_available']:
        print_status("System check passed! Ready to run platform", "success")
        
        if dep_results['missing_optional']:
            print_status(
                f"Note: {len(dep_results['missing_optional'])} optional packages not installed. "
                "Some features may be limited.",
                "info"
            )
        
        return True
    else:
        print_status("System check failed", "error")
        return False


def show_installation_guide():
    """Show installation guide for missing dependencies."""
    print(f"{Colors.BOLD}Installation Guide:{Colors.END}")
    print()
    
    print("1. Install required packages:")
    print("   pip install numpy scipy scikit-learn psutil")
    print()
    
    print("2. Install optional packages for enhanced features:")
    print("   pip install numba yaml toml aiohttp fastapi uvicorn")
    print()
    
    print("3. For AI/LLM integration:")
    print("   pip install openai transformers torch")
    print()
    
    print("4. Run system check:")
    print("   python run.py --check")
    print()


def interactive_mode():
    """Interactive mode for selecting startup options."""
    print_header()
    
    print("Welcome to the AI Historical Simulation Platform!")
    print()
    print("What would you like to do?")
    print()
    print("1. 💬 Interactive Chat - Chat with historical figures")
    print("2. 🎬 Demo Mode - Run pre-configured scenarios")
    print("3. 📊 System Status - Check platform status")
    print("4. 📚 List Figures - See available historical figures")
    print("5. 🔧 System Check - Validate system requirements")
    print("6. 📦 Install Dependencies - Install missing packages")
    print("7. ❓ Help - Show help information")
    print("8. 🚪 Exit")
    print()
    
    while True:
        try:
            choice = input("Select option (1-8): ").strip()
            
            if choice == '1':
                return ['app.py', 'chat']
            elif choice == '2':
                return ['app.py', 'demo']
            elif choice == '3':
                return ['app.py', 'status']
            elif choice == '4':
                return ['app.py', 'figures']
            elif choice == '5':
                run_system_check()
                input("\nPress Enter to continue...")
                continue
            elif choice == '6':
                dep_results = check_dependencies()
                if dep_results['missing_required']:
                    install_missing_dependencies(dep_results)
                else:
                    print_status("No missing required dependencies", "success")
                input("\nPress Enter to continue...")
                continue
            elif choice == '7':
                return ['app.py', 'help']
            elif choice == '8':
                print_status("Goodbye!", "info")
                return None
            else:
                print_status("Invalid choice. Please select 1-8.", "warning")
                
        except KeyboardInterrupt:
            print_status("\nGoodbye!", "info")
            return None
        except Exception as e:
            print_status(f"Error: {e}", "error")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="AI Historical Simulation Platform Startup Script"
    )
    parser.add_argument('--chat', action='store_true', help='Start in chat mode')
    parser.add_argument('--demo', action='store_true', help='Start in demo mode')
    parser.add_argument('--status', action='store_true', help='Show platform status')
    parser.add_argument('--figures', action='store_true', help='List available figures')
    parser.add_argument('--check', action='store_true', help='Run system check')
    parser.add_argument('--install-deps', action='store_true', help='Install missing dependencies')
    parser.add_argument('--help-install', action='store_true', help='Show installation guide')
    parser.add_argument('--no-check', action='store_true', help='Skip system check')
    
    args = parser.parse_args()
    
    # Handle specific arguments
    if args.help_install:
        show_installation_guide()
        return
    
    if args.check:
        success = run_system_check()
        sys.exit(0 if success else 1)
    
    if args.install_deps:
        dep_results = check_dependencies()
        if install_missing_dependencies(dep_results):
            print_status("Installation completed successfully", "success")
        else:
            print_status("Installation failed", "error")
            sys.exit(1)
        return
    
    # Determine mode
    app_args = []
    if args.chat:
        app_args = ['app.py', 'chat']
    elif args.demo:
        app_args = ['app.py', 'demo']
    elif args.status:
        app_args = ['app.py', 'status']
    elif args.figures:
        app_args = ['app.py', 'figures']
    else:
        # Interactive mode
        app_args = interactive_mode()
        if app_args is None:
            return
    
    # Run system check unless explicitly skipped
    if not args.no_check:
        print_status("Running quick system check...", "working")
        if not check_python_version():
            sys.exit(1)
        
        # Quick dependency check
        dep_results = check_dependencies()
        if not dep_results['all_required_available']:
            print_status("Missing required dependencies. Run with --install-deps to install.", "error")
            show_installation_guide()
            sys.exit(1)
        
        print_status("System check passed", "success")
        print()
    
    # Launch the application
    try:
        print_status(f"Starting application: {' '.join(app_args[1:])}", "working")
        print()
        
        # Import and run app
        os.execv(sys.executable, [sys.executable] + app_args)
        
    except Exception as e:
        print_status(f"Failed to start application: {e}", "error")
        sys.exit(1)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print_status("\nGoodbye!", "info")
    except Exception as e:
        print_status(f"Fatal error: {e}", "error")
        sys.exit(1)